#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
STM32 Firmware Side-Channel Attack Analysis Tool (External Interface)
Calls side_channel_attacks module to perform attack analysis
Supports: multiple algorithms, leakage models, plaintext loading, byte-range attacks, and other advanced options
"""

import sys
import os
import numpy as np
import argparse
import re
import json
from datetime import datetime
from typing import Optional, Dict, Any, List
from pathlib import Path

# Add side_channel_attacks to path
sys.path.insert(0, os.path.join(os.path.dirname(__file__), 'side_channel_attacks'))

# Import attack manager
try:
    from attack_manager import SideChannelAttackManager
    from base_side_channel_attack import AttackResult
except ModuleNotFoundError as e:
    print(f"Error: Cannot import side-channel attack module: {e}")
    sys.exit(1)


class FirmwareTraceAnalyzer:
    """STM32 Firmware Trace Analyzer (External Interface Wrapper)"""
    
    def __init__(self, trace_file: str, algorithm: str = 'aes',
                 plaintext_file: Optional[str] = None,
                 leakage_model: str = 'hamming_weight',
                 log_level: str = 'INFO'):
        """
        Initialize analyzer
        
        Args:
            trace_file: Trace file in NPZ format
            algorithm: Cryptographic algorithm type
            plaintext_file: Plaintext data file path (optional)
            leakage_model: Leakage model type
            log_level: Logging level
        """
        self.trace_file = trace_file
        self.algorithm = algorithm.lower()
        self.plaintext_file = plaintext_file
        self.leakage_model = leakage_model
        
        # Create attack manager
        self.manager = SideChannelAttackManager(log_level=log_level)
        
        # Store loaded data
        self.loaded_plaintexts = None
        self.original_data = None
    
    def load_plaintexts_from_header(self, header_file: str) -> Optional[np.ndarray]:
        """Load plaintext data from C header file"""
        try:
            with open(header_file, 'r') as f:
                content = f.read()
            
            # 解析NUM_TRACES
            num_match = re.search(r'#define\s+NUM_TRACES\s+(\d+)', content)
            if num_match:
                num_traces = int(num_match.group(1))
            else:
                print(f"  警告: 无法从 {header_file} 解析NUM_TRACES")
                return None
            
            # 解析test_data数组
            pattern = r'static\s+const\s+uint8_t\s+test_data\[.*?\]\[16\]\s*=\s*\{([^}]+)\}'
            match = re.search(pattern, content, re.DOTALL)
            
            if not match:
                print(f"  警告: 无法从 {header_file} 解析test_data")
                return None
            
            data_str = match.group(1)
            hex_values = re.findall(r'0x([0-9a-fA-F]{2})', data_str)
            
            if len(hex_values) < num_traces * 16:
                print(f"  警告: 头文件中数据不足")
                return None
            
            # 转换为numpy数组
            plaintexts = np.zeros((num_traces, 16), dtype=np.uint8)
            for i in range(num_traces):
                for j in range(16):
                    idx = i * 16 + j
                    if idx < len(hex_values):
                        plaintexts[i, j] = int(hex_values[idx], 16)
            
            print(f"  [OK] 从 {header_file} 加载了 {num_traces} 个明文")
            return plaintexts
            
        except Exception as e:
            print(f"  警告: 加载明文文件失败: {e}")
            return None
    
    def prepare_npz_with_plaintexts(self) -> str:
        """准备包含明文的临时NPZ文件"""
        try:
            # 加载原始NPZ数据
            print(f"\n[1] 加载轨迹文件: {self.trace_file}")
            data = np.load(self.trace_file, allow_pickle=True)
            self.original_data = dict(data)
            
            # 显示文件信息
            print(f"[OK] 轨迹文件加载成功")
            print(f"  文件内容: {list(data.keys())}")
            if 'num_traces' in data:
                print(f"  轨迹数量: {data['num_traces']}")
            if 'operation' in data:
                print(f"  操作类型: {data['operation']}")
            
            # 获取CPU轨迹（支持多种命名）
            cpu_traces = data.get('cpu_traces', data.get('traces', data.get('power_traces')))
            if cpu_traces is None:
                print("错误: 未找到轨迹数据 (cpu_traces, traces, 或 power_traces)")
                return None
            
            num_traces = len(cpu_traces)
            print(f"  CPU轨迹形状: {cpu_traces.shape}")
            
            # 准备明文数据
            plaintexts = None
            
            # 优先级1: NPZ文件中的明文
            if 'plaintexts' in data:
                plaintexts = data['plaintexts']
                print(f"  [OK] 使用NPZ文件中的明文")
            
            # 优先级2: 外部明文文件
            elif self.plaintext_file:
                if os.path.exists(self.plaintext_file):
                    if self.plaintext_file.endswith('.h'):
                        plaintexts = self.load_plaintexts_from_header(self.plaintext_file)
                    elif self.plaintext_file.endswith('.npz'):
                        pt_data = np.load(self.plaintext_file)
                        if 'plaintexts' in pt_data:
                            plaintexts = pt_data['plaintexts'][:num_traces]
                            print(f"  [OK] 从外部NPZ加载明文")
                else:
                    print(f"  警告: 明文文件不存在: {self.plaintext_file}")
            
            # 最后: 生成随机明文
            if plaintexts is None:
                print("  警告: 未找到明文数据，生成随机明文")
                print("  提示: 使用 --plaintexts 参数加载真实明文以提高攻击成功率")
                plaintexts = np.random.randint(0, 256, (num_traces, 16), dtype=np.uint8)
            
            # 确保明文数量匹配
            if len(plaintexts) != num_traces:
                print(f"  警告: 明文数量不匹配，调整为 {num_traces}")
                if len(plaintexts) > num_traces:
                    plaintexts = plaintexts[:num_traces]
                else:
                    extra = num_traces - len(plaintexts)
                    plaintexts = np.vstack([
                        plaintexts,
                        np.random.randint(0, 256, (extra, 16), dtype=np.uint8)
                    ])
            
            self.loaded_plaintexts = plaintexts
            
            # 创建临时NPZ文件（如果需要添加明文）
            temp_file = self.trace_file
            if 'plaintexts' not in data or self.plaintext_file:
                # 需要创建新的NPZ文件
                temp_file = self.trace_file.replace('.npz', '_with_plaintexts.npz')
                
                # 准备保存数据
                save_data = {
                    'power_traces': cpu_traces,
                    'traces': cpu_traces,  # 兼容性
                    'plaintexts': plaintexts,
                }
                
                # 处理密钥（确保有keys字段）
                if 'keys' in data:
                    save_data['keys'] = data['keys']
                elif 'key' in data:
                    # 将单个密钥扩展为每条轨迹的密钥数组
                    single_key = data['key']
                    save_data['keys'] = np.tile(single_key, (num_traces, 1))
                    save_data['key'] = single_key  # 也保留原始key
                
                # 复制其他字段
                for key in ['ciphertexts', 'operation', 'num_traces']:
                    if key in data:
                        save_data[key] = data[key]
                
                # 保存临时文件
                np.savez_compressed(temp_file, **save_data)
                print(f"  [OK] 创建临时文件: {temp_file}")
            
            return temp_file
            
        except Exception as e:
            print(f"错误: 准备数据失败: {e}")
            import traceback
            traceback.print_exc()
            return None
    
    def _analyze_fuzzy_match(self, byte_results: List, correct_key: np.ndarray,
                            byte_start: int, byte_end: int, top_n: int) -> Dict[str, Any]:
        """
        分析模糊匹配：检查正确密钥字节在top_n候选中的排名
        
        Args:
            byte_results: 每个字节的攻击结果
            correct_key: 正确的密钥
            byte_start: 起始字节
            byte_end: 结束字节
            top_n: 候选数量
        
        Returns:
            模糊匹配统计信息
        """
        stats = {
            'in_top_n': 0,  # 在top_n中的字节数
            'exact_match': 0,  # 排名第一的字节数
            'rankings': [],  # 每个字节的排名
            'avg_rank': 0.0,  # 平均排名
            'total_bytes': byte_end - byte_start
        }
        
        for idx, (byte_idx, result) in enumerate(zip(range(byte_start, byte_end), byte_results)):
            if result is None or not hasattr(result, 'top_candidates'):
                stats['rankings'].append(None)
                continue
            
            correct_byte_value = int(correct_key[byte_idx])
            
            # 查找正确密钥字节的排名
            rank = result.get_rank_of_hypothesis(correct_byte_value)
            
            # 调试：显示前几个候选和正确值（已关闭）
            # if idx < 2:  # 只显示前2个字节
            #     print(f"  [DEBUG] Byte {byte_idx}: correct={correct_byte_value}, top3={result.top_candidates[:3]}, rank={rank}")
            
            if rank is not None:
                stats['rankings'].append(rank)
                
                if rank == 1:
                    stats['exact_match'] += 1
                
                if rank <= top_n:
                    stats['in_top_n'] += 1
            else:
                stats['rankings'].append(None)
        
        # 计算平均排名（仅计算有排名的字节）
        valid_ranks = [r for r in stats['rankings'] if r is not None]
        if valid_ranks:
            stats['avg_rank'] = sum(valid_ranks) / len(valid_ranks)
        
        return stats
    
    def _print_fuzzy_match_stats(self, fuzzy_stats: Dict[str, Any], top_n: int):
        """打印模糊匹配统计信息"""
        print(f"\n  模糊匹配分析 (Top-{top_n}):")
        print(f"    - 排名第1 (精确): {fuzzy_stats['exact_match']}/{fuzzy_stats['total_bytes']} 字节")
        print(f"    - 在Top-{top_n}内: {fuzzy_stats['in_top_n']}/{fuzzy_stats['total_bytes']} 字节")
        
        if fuzzy_stats['avg_rank'] > 0:
            print(f"    - 平均排名: {fuzzy_stats['avg_rank']:.2f}")
        
        # 显示每个字节的详细排名
        print(f"    - 字节排名详情: ", end="")
        rankings_display = []
        for i, rank in enumerate(fuzzy_stats['rankings']):
            if rank is None:
                rankings_display.append("N/A")
            elif rank <= top_n:
                rankings_display.append(f"#{rank}")
            else:
                rankings_display.append(f">{top_n}")
        print(", ".join(rankings_display))
        
        # 评估攻击效果
        success_rate = fuzzy_stats['in_top_n'] / fuzzy_stats['total_bytes'] if fuzzy_stats['total_bytes'] > 0 else 0
        print(f"\n  攻击效果评估:")
        if success_rate >= 0.8:
            print(f"    状态: 优秀 ({success_rate*100:.1f}% 在Top-{top_n}内)")
            print(f"    结论: 攻击非常有效，密钥可通过枚举Top-{top_n}候选恢复")
        elif success_rate >= 0.5:
            print(f"    状态: 良好 ({success_rate*100:.1f}% 在Top-{top_n}内)")
            print(f"    结论: 攻击有效，需要适度的枚举工作")
        elif success_rate >= 0.3:
            print(f"    状态: 中等 ({success_rate*100:.1f}% 在Top-{top_n}内)")
            print(f"    结论: 攻击部分有效，需要增加轨迹数量或优化攻击参数")
        else:
            print(f"    状态: 较差 ({success_rate*100:.1f}% 在Top-{top_n}内)")
            print(f"    结论: 攻击效果不佳，建议检查轨迹质量或更换攻击方法")
    
    def _generate_detailed_report(self, attack_results: Dict[str, Dict], 
                                  byte_range: tuple, top_n: int,
                                  target_operation: str = None) -> Dict[str, Any]:
        """
        生成详细的JSON攻击报告
        
        Args:
            attack_results: 所有攻击的结果字典
            byte_range: 攻击的字节范围
            top_n: 候选数量
            target_operation: 目标操作
        
        Returns:
            详细报告字典
        """
        report = {
            'metadata': {
                'timestamp': datetime.now().isoformat(),
                'trace_file': self.trace_file,
                'algorithm': self.algorithm,
                'leakage_model': self.leakage_model,
                'target_operation': target_operation,
                'byte_range': {
                    'start': byte_range[0],
                    'end': byte_range[1],
                    'total_bytes': byte_range[1] - byte_range[0]
                },
                'top_n': top_n
            },
            'trace_info': {},
            'attacks': []
        }
        
        # 添加轨迹信息
        if self.manager.traces is not None:
            report['trace_info'] = {
                'num_traces': int(self.manager.traces.shape[0]),
                'trace_length': int(self.manager.traces.shape[1]),
                'has_keys': self.manager.keys is not None,
                'has_plaintexts': self.manager.plaintexts is not None
            }
        
        # 添加每种攻击的详细结果
        for attack_name, result in attack_results.items():
            attack_report = {
                'attack_type': attack_name,
                'success': result.get('success', False),
                'avg_confidence': float(result.get('avg_confidence', 0.0)),
                'total_time': float(result.get('total_time', 0.0)),
                'recovered_key': result.get('recovered_key', bytearray(16)).hex(),
                'exact_match': {
                    'correct_bytes': int(result.get('correct_bytes', 0)),
                    'total_bytes': byte_range[1] - byte_range[0]
                }
            }
            
            # 添加正确密钥（如果有）
            if self.manager.keys is not None:
                correct_key = None
                if isinstance(self.manager.keys, np.ndarray):
                    if self.manager.keys.ndim == 1:
                        correct_key = self.manager.keys
                    elif self.manager.keys.ndim > 1:
                        correct_key = self.manager.keys[0]
                
                if correct_key is not None:
                    attack_report['correct_key'] = bytes(correct_key).hex()
            
            # 添加模糊匹配统计
            fuzzy_stats = result.get('fuzzy_match', {})
            if fuzzy_stats:
                attack_report['fuzzy_match'] = {
                    'exact_match_count': int(fuzzy_stats.get('exact_match', 0)),
                    'in_top_n_count': int(fuzzy_stats.get('in_top_n', 0)),
                    'avg_rank': float(fuzzy_stats.get('avg_rank', 0.0)),
                    'rankings': [int(r) if r is not None else None for r in fuzzy_stats.get('rankings', [])],
                    'success_rate': float(fuzzy_stats.get('in_top_n', 0) / fuzzy_stats.get('total_bytes', 1))
                }
            
            # 添加每个字节的详细结果
            if 'byte_details' in result:
                attack_report['byte_details'] = result['byte_details']
            
            report['attacks'].append(attack_report)
        
        return report
    
    def _save_report_to_file(self, report: Dict[str, Any], filename: str):
        """保存报告到JSON文件"""
        try:
            with open(filename, 'w', encoding='utf-8') as f:
                json.dump(report, f, indent=2, ensure_ascii=False)
            print(f"\n详细报告已保存: {filename}")
        except Exception as e:
            print(f"\n警告: 保存报告失败: {e}")
    
    def recover_full_key(self, attack_type: str, byte_range: tuple,
                        target_operation: str = None, verbose: bool = False,
                        top_n: int = 10) -> Dict[str, Any]:
        """
        恢复完整密钥（攻击多个字节）
        
        Args:
            attack_type: 攻击类型
            byte_range: 要攻击的字节范围 (start, end)
            target_operation: 目标操作（算法特定攻击）
            verbose: 是否显示详细输出
            top_n: 保存前n个候选用于模糊匹配
        
        Returns:
            恢复结果字典
        """
        byte_start, byte_end = byte_range
        print(f"\n[3] 执行 {attack_type.upper()} 攻击")
        print(f"  算法: {self.algorithm.upper()}")
        print(f"  泄露模型: {self.leakage_model}")
        print(f"  字节范围: {byte_start}-{byte_end-1}")
        
        recovered_key = bytearray(16)
        confidences = []
        execution_times = []
        byte_results = []  # 保存每个字节的详细结果用于模糊匹配
        byte_details = []  # 保存每个字节的详细信息用于报告
        
        for byte_idx in range(byte_start, byte_end):
            if verbose:
                print(f"    攻击字节 {byte_idx}...")
            
            try:
                # 根据攻击类型调用不同方法
                if attack_type in ['cpa', 'dpa', 'spa', 'template']:
                    # 基础攻击
                    result = self.manager.run_basic_attack(
                        attack_type,
                        target_byte=byte_idx,
                        target_operation=target_operation if attack_type == 'template' else None,
                        leakage_model=self.leakage_model if attack_type == 'template' else None,
                        top_n=top_n
                    )
                elif attack_type.startswith('algo_'):
                    # 算法特定攻击
                    actual_type = attack_type.replace('algo_', '')
                    result = self.manager.run_algorithm_specific_attack(
                        actual_type,
                        target_operation=target_operation or 'first_round',
                        target_byte=byte_idx,
                        top_n=top_n
                    )
                else:
                    print(f"  警告: 未知攻击类型 {attack_type}")
                    continue
                
                # 提取结果
                if result and result.success and hasattr(result, 'best_hypothesis'):
                    recovered_key[byte_idx] = result.best_hypothesis
                    confidences.append(result.confidence)
                    execution_times.append(result.execution_time)
                    byte_results.append(result)  # 保存完整结果用于模糊匹配
                    
                    # 收集字节详细信息用于报告
                    byte_detail = {
                        'byte_index': int(byte_idx),
                        'success': True,
                        'best_hypothesis': int(result.best_hypothesis),
                        'confidence': float(result.confidence),
                        'execution_time': float(result.execution_time),
                        'top_candidates': [(int(h), float(s)) for h, s in (result.top_candidates[:10] if hasattr(result, 'top_candidates') else [])]
                    }
                    byte_details.append(byte_detail)
                    
                    if verbose:
                        print(f"      结果: 0x{result.best_hypothesis:02x} " +
                              f"(置信度: {result.confidence:.4f}, " +
                              f"耗时: {result.execution_time:.2f}s)")
                else:
                    print(f"  警告: 字节 {byte_idx} 攻击失败")
                    recovered_key[byte_idx] = 0
                    byte_results.append(None)
                    confidences.append(0.0)
                    execution_times.append(0.0)
                    
                    byte_detail = {
                        'byte_index': int(byte_idx),
                        'success': False,
                        'best_hypothesis': 0,
                        'confidence': 0.0,
                        'execution_time': 0.0,
                        'top_candidates': []
                    }
                    byte_details.append(byte_detail)
                    
            except Exception as e:
                print(f"  错误: 字节 {byte_idx} 攻击异常: {e}")
                recovered_key[byte_idx] = 0
                byte_results.append(None)
                confidences.append(0.0)
                execution_times.append(0.0)
                
                byte_detail = {
                    'byte_index': int(byte_idx),
                    'success': False,
                    'error': str(e),
                    'best_hypothesis': 0,
                    'confidence': 0.0,
                    'execution_time': 0.0,
                    'top_candidates': []
                }
                byte_details.append(byte_detail)
        
        # 计算统计信息
        avg_confidence = np.mean(confidences) if confidences else 0.0
        total_time = sum(execution_times)
        
        # 与正确密钥对比（精确匹配和模糊匹配）
        correct_key = None
        correct_bytes = 0
        fuzzy_match_stats = {}
        
        if self.manager.keys is not None:
            # 处理不同的密钥格式
            if isinstance(self.manager.keys, np.ndarray):
                if self.manager.keys.ndim == 1:
                    correct_key = self.manager.keys
                elif self.manager.keys.ndim > 1:
                    correct_key = self.manager.keys[0]
            
            if correct_key is not None and len(correct_key) > 0:
                # 精确匹配
                correct_bytes = sum(1 for i in range(byte_start, byte_end)
                                  if recovered_key[i] == correct_key[i])
                
                # 模糊匹配：检查正确密钥字节在top_n中的排名
                fuzzy_match_stats = self._analyze_fuzzy_match(
                    byte_results, correct_key, byte_start, byte_end, top_n
                )
        
        # 显示结果
        print(f"\n{attack_type.upper()} 攻击完成!")
        print(f"  恢复密钥: {recovered_key.hex()}")
        if correct_key is not None:
            print(f"  正确密钥: {bytes(correct_key).hex()}")
            print(f"  精确匹配: {correct_bytes}/{byte_end - byte_start} 字节")
            
            # 显示模糊匹配统计
            if fuzzy_match_stats:
                self._print_fuzzy_match_stats(fuzzy_match_stats, top_n)
        
        print(f"  平均置信度: {avg_confidence:.4f}")
        print(f"  总耗时: {total_time:.2f}秒")
        
        return {
            'success': correct_bytes == (byte_end - byte_start),
            'recovered_key': recovered_key,
            'correct_bytes': correct_bytes,
            'avg_confidence': avg_confidence,
            'total_time': total_time,
            'fuzzy_match': fuzzy_match_stats,
            'byte_details': byte_details
        }
    
    def analyze(self, attack_types: List[str] = None,
               byte_range: tuple = None,
               target_operation: str = 'first_round',
               verbose: bool = False,
               save_result: Optional[str] = None,
               top_n: int = 10,
               report_file: Optional[str] = None) -> bool:
        """
        执行完整分析流程
        
        Args:
            attack_types: 要执行的攻击类型列表
            byte_range: 要攻击的字节范围 (start, end)
            target_operation: 算法特定攻击的目标操作
            verbose: 是否显示详细输出
            save_result: 保存结果的文件路径
            top_n: 保存前n个候选用于模糊匹配
            report_file: 保存详细JSON报告的文件路径
        
        Returns:
            是否成功
        """
        if attack_types is None:
            attack_types = ['cpa']
        
        if byte_range is None:
            byte_range = (0, 16)
        
        print("=" * 70)
        print("  STM32固件侧信道攻击分析工具")
        print("=" * 70)
        
        # 步骤1: 准备数据（包含明文）
        npz_file = self.prepare_npz_with_plaintexts()
        if npz_file is None:
            return False
        
        # 步骤2: 使用攻击管理器加载数据
        print(f"\n[2] 初始化攻击管理器...")
        if not self.manager.load_data(npz_file, algorithm=self.algorithm):
            print("错误: 数据加载失败")
            return False
        
        # 步骤3: 执行攻击
        results = {}
        for attack_type in attack_types:
            result = self.recover_full_key(
                attack_type=attack_type,
                byte_range=byte_range,
                target_operation=target_operation,
                verbose=verbose,
                top_n=top_n
            )
            results[attack_type] = result
        
        # 步骤4: 总结
        print("\n" + "=" * 70)
        print("  攻击结果总结")
        print("=" * 70)
        
        for attack_name, result in results.items():
            if result and result['success']:
                print(f"[OK] {attack_name.upper()}: 成功 " +
                      f"(置信度: {result['confidence']:.4f}, " +
                      f"正确字节: {result.get('correct_bytes', '?')}/{byte_range[1]-byte_range[0]})")
            else:
                print(f"[FAIL] {attack_name.upper()}: 失败")
        
        print("=" * 70)
        
        # 步骤5: 保存结果（可选）
        if save_result:
            self.save_results(results, save_result, byte_range)
        
        # 步骤6: 生成详细报告（可选）
        if report_file:
            report = self._generate_detailed_report(results, byte_range, top_n, target_operation)
            self._save_report_to_file(report, report_file)
        
        # 清理临时文件
        if npz_file != self.trace_file and os.path.exists(npz_file):
            os.remove(npz_file)
            if verbose:
                print(f"  清理临时文件: {npz_file}")
        
        return True
    
    def save_results(self, results: Dict[str, Any], filename: str, byte_range: tuple):
        """保存攻击结果到文件"""
        try:
            print(f"\n保存结果到: {filename}")
            
            save_data = {
                'algorithm': self.algorithm,
                'leakage_model': self.leakage_model,
                'byte_range_start': byte_range[0],
                'byte_range_end': byte_range[1],
                'num_traces': len(self.manager.traces) if self.manager.traces is not None else 0,
                'trace_length': self.manager.traces.shape[1] if self.manager.traces is not None else 0,
            }
            
            # 保存各个攻击的结果
            for attack_name, result in results.items():
                if result:
                    save_data[f'{attack_name}_recovered_key'] = result['recovered_key'].hex()
                    save_data[f'{attack_name}_confidence'] = float(result['confidence'])
                    save_data[f'{attack_name}_execution_time'] = float(result['execution_time'])
                    if 'correct_bytes' in result:
                        save_data[f'{attack_name}_correct_bytes'] = int(result['correct_bytes'])
            
            # 保存正确密钥
            if self.manager.keys is not None and len(self.manager.keys) > 0:
                save_data['true_key'] = bytes(self.manager.keys[0]).hex()
            
            np.savez(filename, **save_data)
            print(f"[OK] 结果已保存")
            
        except Exception as e:
            print(f"警告: 保存结果失败: {e}")


def main():
    parser = argparse.ArgumentParser(
        description='STM32固件侧信道攻击分析工具 - 调用side_channel_attacks模块',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
示例:
  # 基本CPA攻击
  analyze_firmware_traces.py traces.npz --attacks cpa
  
  # 指定算法和泄露模型
  analyze_firmware_traces.py traces.npz --algorithm aes --leakage-model hamming_weight
  
  # 从头文件加载真实明文
  analyze_firmware_traces.py traces.npz --plaintexts bin/aes_test_data.h --attacks cpa
  
  # 执行多种攻击
  analyze_firmware_traces.py traces.npz --attacks cpa dpa algo_cpa
  
  # 只攻击特定字节范围（快速测试）
  analyze_firmware_traces.py traces.npz --byte-range 0,4 --attacks cpa --verbose
  
  # 模糊匹配分析（评估攻击效果）
  analyze_firmware_traces.py traces.npz --attacks cpa --top-n 50 --byte-range 0,8
  analyze_firmware_traces.py traces.npz --attacks cpa --top-n 100  # 更大的候选范围
  
  # 生成详细JSON报告
  analyze_firmware_traces.py traces.npz --attacks cpa --report attack_report.json
  analyze_firmware_traces.py traces.npz --attacks cpa dpa --top-n 50 --report detailed_report.json
        """)
    
    # 必需参数
    parser.add_argument('trace_file', type=str,
                       help='轨迹文件路径 (.npz格式)')
    
    # 攻击选项
    parser.add_argument('--attacks', type=str, nargs='+',
                       default=['cpa'],
                       choices=['cpa', 'dpa', 'spa', 'template', 'algo_cpa', 'algo_dpa'],
                       help='要执行的攻击类型 (默认: cpa)')
    
    # 算法配置
    parser.add_argument('--algorithm', '-a', type=str, default='aes',
                       choices=['aes', 'sm4', 'sm2', 'rsa', 'ecc', 'dilithium'],
                       help='密码算法类型 (默认: aes)')
    
    parser.add_argument('--leakage-model', '-l', type=str, default='hamming_weight',
                       choices=['hamming_weight', 'hamming_distance', 'identity'],
                       help='泄露模型类型 (默认: hamming_weight)')
    
    # 数据源
    parser.add_argument('--plaintexts', '-p', type=str, default=None,
                       metavar='FILE',
                       help='明文数据文件 (.h头文件或.npz)。如不指定则生成随机明文')
    
    # 攻击范围
    def byte_range_type(s):
        """解析字节范围参数"""
        try:
            start, end = map(int, s.split(','))
            if start < 0 or end < 0:
                raise argparse.ArgumentTypeError("字节索引必须为非负数")
            if start >= end:
                raise argparse.ArgumentTypeError("起始字节必须小于结束字节")
            if end > 16:
                raise argparse.ArgumentTypeError("结束字节不能超过16")
            return (start, end)
        except ValueError:
            raise argparse.ArgumentTypeError(
                f"字节范围格式错误: '{s}'. 请使用 'START,END' 格式，例如: 0,4"
            )
    
    parser.add_argument('--byte-range', type=byte_range_type, default=None,
                       metavar='START,END',
                       help='只攻击指定字节范围 [START,END)，例如: 0,4 表示攻击字节0-3 (默认: 0,16即全部字节)')
    
    parser.add_argument('--target-operation', type=str, default='first_round',
                       choices=['first_round', 'last_round', 'sbox_output', 'key_schedule'],
                       help='算法特定攻击的目标操作 (默认: first_round)')
    
    # 输出选项
    parser.add_argument('--verbose', '-v', action='store_true',
                       help='显示详细输出')
    
    parser.add_argument('--save-result', type=str, default=None,
                       metavar='FILE',
                       help='保存攻击结果到文件')
    
    parser.add_argument('--top-n', type=int, default=10,
                       metavar='N',
                       help='保存前N个候选用于模糊匹配分析 (默认: 10)')
    
    parser.add_argument('--report', type=str, default=None,
                       metavar='FILE',
                       help='保存详细JSON格式攻击报告到文件')
    
    parser.add_argument('--log-level', type=str, default='INFO',
                       choices=['DEBUG', 'INFO', 'WARNING', 'ERROR'],
                       help='日志级别 (默认: INFO)')
    
    args = parser.parse_args()
    
    # 创建分析器
    analyzer = FirmwareTraceAnalyzer(
        trace_file=args.trace_file,
        algorithm=args.algorithm,
        plaintext_file=args.plaintexts,
        leakage_model=args.leakage_model,
        log_level=args.log_level
    )
    
    # 执行分析
    success = analyzer.analyze(
        attack_types=args.attacks,
        byte_range=args.byte_range,  # 已经是tuple或None
        target_operation=args.target_operation,
        verbose=args.verbose,
        save_result=args.save_result,
        top_n=args.top_n,
        report_file=args.report
    )
    
    return 0 if success else 1


if __name__ == '__main__':
    exit(main())
